"""
文档处理器模块
负责文档解析、分割和元数据提取
"""
import os
import hashlib
from typing import List, Dict, Any, Optional
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document

class DocumentProcessor:
    """
    文档处理器，负责文档的解析、分割和元数据提取
    """
    
    def __init__(self):
        """
        初始化文档处理器
        """
        # 使用递归字符文本分割器，减小块大小以提高搜索精度
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,         # 减小块大小，使每个块更加聚焦
            chunk_overlap=50,       # 相应调整重叠部分
            length_function=len,    # 长度计算函数
            # 优先按段落和标题分割
            separators=["\n## ", "\n### ", "\n\n", "\n", " ", ""]  # 优先按标题级别分割
        )
    
    def split_text(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> List[Document]:
        """
        将文本分割成多个文档块
        
        Args:
            text: 要分割的文本内容
            metadata: 文档元数据
            
        Returns:
            分割后的文档块列表
        """
        if not metadata:
            metadata = {}
            
        # 首先尝试特殊处理标题和内容，确保重要章节不被分割
        # 这对于结构化文档特别有用，能确保像"RAG的优势"这样的章节保持完整
        enhanced_chunks = []
        
        # 按标题分割的简单实现
        sections = text.split('\n## ')
        for i, section in enumerate(sections):
            if i == 0:
                # 第一部分可能是文档标题和简介
                if len(section) > 0:
                    enhanced_chunks.append(section)
            else:
                # 其他部分是带标题的章节
                section_text = '## ' + section
                if len(section_text) > 500:
                    # 如果章节太长，再使用text_splitter分割
                    sub_chunks = self.text_splitter.split_text(section_text)
                    enhanced_chunks.extend(sub_chunks)
                else:
                    # 保持章节完整
                    enhanced_chunks.append(section_text)
        
        # 创建文档对象列表
        documents = []
        for i, chunk in enumerate(enhanced_chunks):
            # 为每个块创建带有索引的元数据
            chunk_metadata = metadata.copy()
            chunk_metadata["chunk_index"] = i
            chunk_metadata["total_chunks"] = len(enhanced_chunks)
            
            # 尝试提取当前块的标题作为元数据
            lines = chunk.split('\n')
            for line in lines:
                if line.startswith('## '):
                    chunk_metadata["section_title"] = line[3:].strip()  # 提取##后的标题文本
                    break
            
            document = Document(page_content=chunk, metadata=chunk_metadata)
            documents.append(document)
            
        return documents
    
    def calculate_file_hash(self, file_path: str) -> Optional[str]:
        """
        计算文件内容的哈希值
        
        Args:
            file_path: 文件路径
            
        Returns:
            文件内容的SHA256哈希值，如果出错则返回None
        """
        try:
            sha256_hash = hashlib.sha256()
            with open(file_path, 'rb') as file:
                # 分块读取文件内容并更新哈希
                for byte_block in iter(lambda: file.read(4096), b""):
                    sha256_hash.update(byte_block)
            return sha256_hash.hexdigest()
        except Exception as e:
            print(f"计算文件哈希值时出错: {str(e)}")
            return None
    
    def process_document(self, file_path: str, check_update: bool = False, existing_hash: Optional[str] = None) -> List[Document]:
        """
        处理单个文档文件
        支持.txt、.md等文本文件格式
        
        Args:
            file_path: 文档文件路径
            check_update: 是否检查文档更新
            existing_hash: 已存在的文档哈希值，用于比较是否更新
            
        Returns:
            处理后的文档块列表，如果文档未更新则返回空列表
        """
        try:
            # 如果需要检查更新且提供了现有哈希值
            if check_update and existing_hash:
                # 计算当前文件的哈希值
                current_hash = self.calculate_file_hash(file_path)
                # 如果哈希值相同，说明文档未更新
                if current_hash == existing_hash:
                    return []
            
            # 获取文件扩展名
            file_extension = os.path.splitext(file_path)[1].lower()
            
            # 提取文件名作为文档标题
            filename = os.path.basename(file_path)
            
            # 计算文件哈希值
            file_hash = self.calculate_file_hash(file_path)
            
            # 创建基本元数据
            metadata = {
                "source": file_path,
                "filename": filename,
                "file_type": file_extension,
                "processed_at": str(os.path.getmtime(file_path)),
                "content_hash": file_hash
            }
            
            # 读取文件内容
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            
            # 分割文本
            return self.split_text(content, metadata)
            
        except Exception as e:
            print(f"处理文档 {file_path} 时出错: {str(e)}")
            return []
    
    def process_documents_in_directory(self, directory_path: str, check_update: bool = False, existing_hashes: Optional[Dict[str, str]] = None) -> List[Document]:
        """
        处理目录中的所有文档文件
        
        Args:
            directory_path: 包含文档的目录路径
            check_update: 是否检查文档更新
            existing_hashes: 已存在的文档哈希值字典，键为文件路径，值为哈希值
            
        Returns:
            处理后的所有文档块列表，仅包含更新的文档
        """
        documents = []
        
        # 支持的文件类型
        supported_extensions = ['.txt', '.md', '.mdx', '.json', '.csv']
        
        # 如果未提供现有哈希值字典，则初始化为空字典
        if existing_hashes is None:
            existing_hashes = {}
        
        try:
            # 遍历目录中的所有文件
            for root, _, files in os.walk(directory_path):
                for file in files:
                    # 检查文件扩展名
                    if any(file.lower().endswith(ext) for ext in supported_extensions):
                        file_path = os.path.join(root, file)
                        # 获取该文件的现有哈希值（如果有）
                        existing_hash = existing_hashes.get(file_path)
                        # 处理文档并添加到列表
                        file_docs = self.process_document(file_path, check_update, existing_hash)
                        documents.extend(file_docs)
            
        except Exception as e:
            print(f"处理目录 {directory_path} 时出错: {str(e)}")
        
        return documents

# 提供默认实例
document_processor = DocumentProcessor()